#include "capture.hpp"
#include <thread>

gentech::CAVPacketQueue::~CAVPacketQueue()
{
	while (!pkt_queue_.empty()) {
		AVPacket &pkt = pkt_queue_.front();
		if (pkt.data) av_free_packet(&pkt);
		pkt_queue_.pop();
	}
}

void gentech::CAVPacketQueue::set_abort(int abort)
{
	std::unique_lock<std::mutex> lck(mtx_);
	abort_ = abort;
}

void gentech::CAVPacketQueue::push(AVPacket &pkt)
{
	std::unique_lock<std::mutex> lck(mtx_);
	pkt_queue_.push(pkt);
	cv_.notify_one();
}

bool gentech::CAVPacketQueue::pop_front(AVPacket &pkt)
{
	if (abort_) return false;

	std::unique_lock<std::mutex> lck(mtx_);
	while (pkt_queue_.empty()) cv_.wait(lck);
	if (abort_) return false;
	pkt = pkt_queue_.front();
	pkt_queue_.pop();
	return true;
}

gentech::CCapture::CCapture()
{
	pFormatCtx_ = NULL;
	videoStreamIndex_ = -1;
	pVideoCodecCtx_ = NULL;
	pVideoFrame_ = NULL;
	pBGRVideoFrame_ = NULL;
	pSwsCtx_ = NULL;

	read_packet_thread_exit_ = true;
	abort_read_packet_thread_ = true;

	av_register_all();
	avformat_network_init();
}

gentech::CCapture::~CCapture()
{
	close();
}

bool gentech::CCapture::init(SCaptureParam &capture_param)
{
	std::string &str_video_path = capture_param.str_video_path_;

	int result = 0;
	char errorInfo[100] = { 0 };
	AVDictionary *opts = 0;
	av_dict_set(&opts, "stimeout", "5000", 0);
	result = avformat_open_input(&pFormatCtx_, str_video_path.c_str(), NULL, &opts);
	if (result != 0) {
		av_strerror(result, errorInfo, sizeof(errorInfo));
		char szlog[200];
		memset(szlog, 0, sizeof(szlog));
		sprintf_s(szlog, sizeof(szlog), "open file: %s format contex failed, error msg: %s.",
			str_video_path.c_str(), errorInfo);
		std::cout << szlog << std::endl;
		close(); return false;
	}
	// Read packets of a media file to get stream information,
	// such as codec type, codec id, and so on..
	result = avformat_find_stream_info(pFormatCtx_, NULL);
	if (result < 0) {
		av_strerror(result, errorInfo, sizeof(errorInfo));
		char szlog[200];
		memset(szlog, 0, sizeof(szlog));
		sprintf_s(szlog, sizeof(szlog), "find file: %s stream info failed, error msg: %s.",
			str_video_path.c_str(), errorInfo);
		std::cout << szlog << std::endl;
		close(); return false;
	}
	for (unsigned i = 0; i < pFormatCtx_->nb_streams && i < 2; ++i) {
		if (pFormatCtx_->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO) {
			videoStreamIndex_ = i;
			break;
		}
	}
	if (videoStreamIndex_ == -1) {
		char szlog[200];
		memset(szlog, 0, sizeof(szlog));
		sprintf_s(szlog, sizeof(szlog), "file: %s, can not find video stream or audio stream.",
			str_video_path.c_str());
		std::cout << szlog << std::endl;
		close(); return false;
	}
	if (videoStreamIndex_ != -1) {
		pVideoCodecCtx_ =
			pFormatCtx_->streams[videoStreamIndex_]->codec;
		// check if there is a supported decoder
		AVCodec *pcodec = avcodec_find_decoder(pVideoCodecCtx_->codec_id);
		if (!pcodec) {
			char szlog[200];
			memset(szlog, 0, sizeof(szlog));
			sprintf_s(szlog, sizeof(szlog), "file: %s, can not find supported decoder.",
				str_video_path.c_str());
			std::cout << szlog << std::endl;
			close(); return false;
		}
		// Initialize the AVCodecContext to use the given AVCodec
		result = avcodec_open2(pVideoCodecCtx_, pcodec, NULL);
		if (result != 0) {
			av_strerror(result, errorInfo, sizeof(errorInfo));
			char szlog[200];
			memset(szlog, 0, sizeof(szlog));
			sprintf_s(szlog, sizeof(szlog), "can not open the codec for file: %s, error msg: %s.",
				str_video_path.c_str(), errorInfo);
			std::cout << szlog << std::endl;
			close(); return false;
		}
	}
	pVideoFrame_ = av_frame_alloc();
	if (!pVideoFrame_) {
		char szlog[200];
		memset(szlog, 0, sizeof(szlog));
		sprintf_s(szlog, sizeof(szlog), "can not allocate video frame.");
		std::cout << szlog << std::endl;
		close(); return false;
	}
	pBGRVideoFrame_ = av_frame_alloc();
	if (!pBGRVideoFrame_) {
		char szlog[200];
		memset(szlog, 0, sizeof(szlog));
		sprintf_s(szlog, sizeof(szlog), "can not allocate video frame.");
		std::cout << szlog << std::endl;
		close(); return false;
	}
	else {
		pBGRVideoFrame_->format = PIX_FMT_BGR24;
		pBGRVideoFrame_->width = pVideoCodecCtx_->width;
		pBGRVideoFrame_->height = pVideoCodecCtx_->height;
		av_frame_get_buffer(pBGRVideoFrame_, 4);
	}
	pSwsCtx_ = sws_getContext(pVideoCodecCtx_->width, pVideoCodecCtx_->height,
		pVideoCodecCtx_->pix_fmt,
		pVideoCodecCtx_->width, pVideoCodecCtx_->height,
		PIX_FMT_BGR24,
		SWS_BILINEAR, NULL, NULL, NULL);
	if (!pSwsCtx_) {
		char szlog[200];
		memset(szlog, 0, sizeof(szlog));
		sprintf_s(szlog, sizeof(szlog), "can not allocate scale context.");
		std::cout << szlog << std::endl;
		close(); return false;
	}
	char szlog[200];
	memset(szlog, 0, sizeof(szlog));
	sprintf_s(szlog, sizeof(szlog), "open video %s succeed.", str_video_path.c_str());
	std::cout << szlog << std::endl;

	abort_read_packet_thread_ = false;
	read_packet_thread_exit_ = false;
	std::thread read_packet_thread(&CCapture::read_packet, this);
	read_packet_thread.detach();

	return true;
}

void gentech::CCapture::read_packet()
{
	char szlog[200];
	memset(szlog, 0, sizeof(szlog));
	sprintf_s(szlog, sizeof(szlog), "read packet thread start.");
	std::cout << szlog << std::endl;

	int result = 0;
	char errorInfo[200] = { 0 };
	while (true) {
		if (abort_read_packet_thread_) break;
		AVPacket packet;
		av_init_packet(&packet);
		result = av_read_frame(pFormatCtx_, &packet);
		if (result < 0) {
			memset(errorInfo, 0, sizeof(errorInfo));
			av_strerror(result, errorInfo, sizeof(errorInfo));
			char szlog[200];
			memset(szlog, 0, sizeof(szlog));
			sprintf_s(szlog, sizeof(szlog), "read_packet thread, av_read_frame failed,"
				"error msg: %s.", errorInfo);
			std::cout << szlog << std::endl;
			break;
		}
		if (packet.stream_index != videoStreamIndex_) {
			av_free_packet(&packet);
			continue;
		}
		av_dup_packet(&packet);
		pkt_queue_.push(packet);
		std::this_thread::yield();
		//std::this_thread::sleep_for(std::chrono::milliseconds(35));
	}
	// read packet thread exit, notify capture function, in case of capture function is waiting
	pkt_queue_.set_abort(1);
	AVPacket pkt;
	pkt.size = 0; pkt.data = 0;
	pkt_queue_.push(pkt);

	// read packet thread exit, notify close function it has exited.
	std::unique_lock<std::mutex> lck(read_packet_thread_exit_mtx_);
	read_packet_thread_exit_ = true;
	read_packet_thread_exit_cv_.notify_one();
	lck.unlock();

	memset(szlog, 0, sizeof(szlog));
	sprintf_s(szlog, sizeof(szlog), "read packet thread end.");
	std::cout << szlog << std::endl;
}

bool gentech::CCapture::capture(cv::Mat &frame)
{
	static cv::Mat static_frame(pBGRVideoFrame_->height, pBGRVideoFrame_->width,
		CV_8UC3, cv::Scalar::all(0));

	AVPacket packet;
	av_init_packet(&packet);
	if (!pkt_queue_.pop_front(packet) || !packet.data) return false;

	int videoFrameFinished = 0;
	int len = avcodec_decode_video2(pVideoCodecCtx_,
		pVideoFrame_, &videoFrameFinished, &packet);
	if (len < 0) {
		av_free_packet(&packet);
		frame.release(); return true;
	}
	if (videoFrameFinished) {
		sws_scale(pSwsCtx_, pVideoFrame_->data, pVideoFrame_->linesize,
			0, pVideoCodecCtx_->height, pBGRVideoFrame_->data, pBGRVideoFrame_->linesize);
		//static_frame.create(pBGRVideoFrame_->height, pBGRVideoFrame_->width, CV_8UC3);
		for (int r = 0; r < pBGRVideoFrame_->height; ++r) {
			memcpy(static_frame.data + r * static_frame.cols * 3,
				pBGRVideoFrame_->data[0] + r * pBGRVideoFrame_->linesize[0],
				sizeof(uint8_t)* static_frame.cols * 3);
		}
		static_frame.copyTo(frame);
	}
	av_free_packet(&packet);
	return true;
}

void gentech::CCapture::close()
{
	// close read packet thread, and wait for it has exited.
	std::unique_lock<std::mutex> lck(read_packet_thread_exit_mtx_);
	abort_read_packet_thread_ = true;
	while (!read_packet_thread_exit_) read_packet_thread_exit_cv_.wait(lck);

	if (pVideoFrame_) {
		av_frame_free(&pVideoFrame_);
		pVideoFrame_ = NULL;
	}
	if (pBGRVideoFrame_) {
		av_frame_free(&pBGRVideoFrame_);
		pVideoFrame_ = NULL;
	}
	if (pVideoCodecCtx_) {
		avcodec_close(pVideoCodecCtx_);
		pVideoCodecCtx_ = NULL;
	}
	if (pFormatCtx_) {
		avformat_close_input(&pFormatCtx_);
	}
	if (pSwsCtx_) {
		sws_freeContext(pSwsCtx_);
		pSwsCtx_ = NULL;
	}
}